import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from itertools import count

# 全局变量
num_states = 3
states = []
state_indices = {}
transition_matrix = None
current_state = None
G = nx.DiGraph()
edge_labels = {}
pos = {}

# 生成随机合法转移矩阵（每行加起来为1）
def generate_random_transition_matrix(n):
    matrix = np.random.rand(n, n)
    row_sums = matrix.sum(axis=1, keepdims=True)
    return matrix / row_sums

# 初始化马尔可夫链
def initialize_chain(n):
    global states, state_indices, transition_matrix, current_state, G, edge_labels, pos
    states = [chr(65 + i) for i in range(n)]  # A, B, C...
    state_indices = {state: idx for idx, state in enumerate(states)}
    transition_matrix = generate_random_transition_matrix(n)

    # 构建有向图
    G.clear()
    for state in states:
        G.add_node(state)

    edge_labels.clear()
    for i, src in enumerate(states):
        for j, dst in enumerate(states):
            prob = transition_matrix[i][j]
            if prob > 0.05:  # 只显示大于某个阈值的概率边
                G.add_edge(src, dst)
                edge_labels[(src, dst)] = f"{prob:.2f}"

    # 布局设置
    pos = nx.circular_layout(G)

    # 初始化当前状态
    current_state = np.random.choice(states)

# 初始化第一次状态链
initialize_chain(num_states)

# 路径记录
path = [current_state]

# 更新函数，用于动画每一帧
def update(frame):
    global current_state, num_states, path

    # 每 50 帧（即 5 秒）更新一次状态数量和转移矩阵
    if frame % 50 == 0:
        new_num_states = np.random.randint(2, 7)  # 随机选择 2~6 个状态
        print(f"\nResetting to {new_num_states} states at frame {frame}")
        initialize_chain(new_num_states)
        path = [current_state]  # 清空路径

    # 如果当前状态无效（比如状态数变了），重置
    if current_state not in states:
        current_state = np.random.choice(states)

    current_idx = state_indices[current_state]
    next_state = np.random.choice(states, p=transition_matrix[current_idx])
    path.append(next_state)
    current_state = next_state

    # 清空并重新绘制
    ax.clear()
    ax.set_title(f"Current State = {current_state} | States Count = {len(states)}")

    # 绘制状态转移图
    nx.draw(G, pos, with_labels=True, node_size=800, node_color='lightblue', font_size=16, ax=ax, arrows=True)

    # 绘制边标签（转移概率）
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax, font_size=10)

    # 高亮当前状态
    nx.draw_networkx_nodes(G, pos, nodelist=[current_state], node_color='orange', node_size=900, ax=ax)

    # 显示路径（最近10步）
    ax.text(0.5, -0.1, 'Path: ' + ' → '.join(path[-10:]), transform=ax.transAxes,
            ha='center', fontsize=10, color='black')

    return []

# 控制动画暂停/继续的函数
def on_key(event):
    if event.key == 'escape':
        ani.event_source.stop()
        print("Animation stopped by pressing ESC.")

# 创建绘图对象
fig, ax = plt.subplots(figsize=(8, 6))
ani = FuncAnimation(fig, update, frames=count(), interval=100, blit=False)

# 绑定键盘事件
fig.canvas.mpl_connect('key_press_event', on_key)

plt.show()

# 输出最终路径
print("\nFinal Markov Chain Path:")
print(" -> ".join(path))
